(*
 * Primitive implementation of directed graphs.
 *
 * author: inky-kinky
 * contact: comamm@yandex.ru
 *
 * This code is distributed under the terms of GPL v.3 license.
 *)

type 'a node = { n : int; data : 'a }
type 'a graph = { mutable nodes : 'a node list; arcs : (int, int list) Hashtbl.t }

exception No_such_node

(* Graph construction. *)
let create n = { nodes = []; arcs = Hashtbl.create n }

let add_arc g n1 n2 = match g.nodes with
  [] -> raise No_such_node
  | h::t -> 
    if (n1 <= h.n) && (n2 <= h.n) then
      try
        let arcs_to = Hashtbl.find g.arcs n1 in
        if not (List.exists (fun x -> x = n2) arcs_to) then
          Hashtbl.replace g.arcs n1 (n2::arcs_to)
        else ()
      with
        Not_found -> Hashtbl.add g.arcs n1 [n2]
    else raise No_such_node

let add_node g x = match g.nodes with
  [] -> g.nodes <- [{ n = 0; data = x}]; 0
  | h::t -> 
    let next = h.n + 1 in 
    g.nodes <- { n = next; data = x}::g.nodes; next

let add_unique_node g x =
  try
    (List.find (fun y -> y.data = x) g.nodes).n 
  with Not_found -> add_node g x

(* Graph statistics. *)
let n_nodes g = match g.nodes with
  [] -> 0
  | h::t -> h.n

let n_arcs g = 
  let counter = ref 0 in
  let count arcs_from arcs_to = 
    counter := !counter + (List.length arcs_to)
  in
  Hashtbl.iter count g.arcs; !counter

(* Graph searching. *)
let exists_node f g = List.exists f g.nodes

let find_node f g = 
  try 
    List.find f g.nodes 
  with Not_found -> raise No_such_node

let neighbors g node =
  try 
    let arcs_to = Hashtbl.find g.arcs node.n in
    List.map (fun x -> List.find (fun y -> y.n = x) g.nodes) arcs_to
  with Not_found -> []

(* Graph random walk. *)
let random_node g = match g.nodes with
  [] -> raise No_such_node
  | h::t -> List.nth g.nodes (Random.int (h.n + 1))

let random_neighbor g node = 
  let all = neighbors g node in
  match all with
    [] -> raise No_such_node
    | h::t -> let l = List.length all in List.nth all (Random.int l)
